import time
import json
import logging
from typing import List, Dict

from fastapi import status, Query
from pydantic import BaseModel
from app.schemas.common import ApiResponse  # 通用响应模型
from app.routers.base_router import BaseRouter
from app.services.multi_collection_query import MultiCollectionQueryService
from app.exception.rag_exception import RAGException
from app.services.parse_file_service import FilesParseService
import re
# 初始化日志
logger = logging.getLogger("SystemQueryRouter")

# --------------------------
# 1. 输入输出模型（Pydantic校验）
# --------------------------
class SystemQueryRequest(BaseModel):
    """用户查询请求参数模型：严格校验输入格式"""
    question: str  # 用户问题（合同法) ”）
    class Config:
        json_schema_extra = {
            "example": {
                "question": "BMI为19，糖尿病，近2个月体重下降4% "
            }
        }

# --------------------------
# 2. 路由实现（遵循BaseRouter规范）
# --------------------------
class SystemQueryRouter(BaseRouter):
    """NRS2002评分查询路由：对外提供HTTP接口"""

    def __init__(self):
        logger.info("Initializing SystemQueryRouter")
        super().__init__()
        # 注册路由
        self.router = self._register_routes()

    def _register_routes(self):
        """注册所有路由并添加Swagger文档注解"""
        # RAG查询接口
        self.router.post(
            "/queryRag",
            response_model=ApiResponse,
            status_code=status.HTTP_200_OK,
            summary="rag查询",
            description="接收用户问题，通过向量检索匹配相关规则，调用Ollama生成结果",
            tags=["合同查询"]
        )(self.sys_contract_query)

        self.router.get(
            "/search",
            response_model=ApiResponse,
            status_code=status.HTTP_200_OK,
            summary="rag查询",
            description="接收用户问题，通过向量检索匹配相关规则，调用Ollama生成结果",
            tags=["合同查询"]
        )(self.search)

        self.router.post(
            "/saveChunk",
            response_model=ApiResponse,
            status_code=status.HTTP_200_OK,
            summary="rag查询",
            description="接收用户问题，通过向量检索匹配相关规则，调用Ollama生成结果",
            tags=["合同查询"]
        )(self.save_chunk)

        return self.router

    async def sys_contract_query(
            self,
            request: SystemQueryRequest
    ) -> ApiResponse:
        """
        核心逻辑：1.校验文件 2.构造ChatBoxApi参数 3.调用RAG查询 4.返回结果
        """
        try:
            # 捕获实例创建异常
            query_service = MultiCollectionQueryService(
                use_llm_rerank=False,
                collections=[
                    "doc_documents_milvus",
                    "docx_documents_milvus",
                    "excel_documents_milvus",
                    "json_documents_milvus",
                    "pdf_documents_milvus",
                    "txt_documents_milvus"
                ]  # 传入实际的集合名称列表
            )
            logger.info(f"服务初始化成功（可用集合数：{len(query_service.vector_stores)}）")
            # 调用查询
            custom_filter = {"file_name LIKE '劳动合同法%'"}
            result = await query_service.query_rag(
                user_question=request.question,
                user_filter=custom_filter
            )
            # 处理结果
            if result["status"] == "error":
                logger.info(f"查询失败：{result['error']['message']}")
            else:
                logger.info(f"LLM响应：{result['response']}")
            return ApiResponse(status=status.HTTP_200_OK,message="success",data=result)
        except RAGException as e:
            logger.info(f"服务初始化失败：[{e.code}] {e.message}")
            logger.error(f"服务初始化失败：[{e.code}] {e.message}", exc_info=True)
        except Exception as e:
            logger.info(f"未知错误：{str(e)}")
            logger.error(f"未知错误：{str(e)}", exc_info=True)

    async def search(
            self,  # 添加self参数，作为类方法的第一个参数
            question: str = Query(..., description="用户查询的问题，例如：劳动合同法中关于试用期的规定有哪些"),
            file_filter: str = Query("劳动合同法%", description="文件名过滤前缀，例如：劳动合同法%")
    ) -> ApiResponse:
        """
        核心逻辑：1.校验参数 2.初始化查询服务 3.调用RAG查询 4.返回结果
        支持GET请求，通过查询参数接收问题和过滤条件
        """
        # 记录开始时间
        start_time = time.time()
        query_service = None
        try:
            # 简单参数校验
            print(f"question:{question}")  # 现在question会正确显示为查询字符串
            if not question.strip():
                return ApiResponse(
                    status=status.HTTP_400_BAD_REQUEST,
                    message="查询问题不能为空"
                )

            # 初始化查询服务（保持不变）
            query_service = MultiCollectionQueryService(
                use_llm_rerank=False,
                collections=[
                    "doc_documents_milvus",
                    "docx_documents_milvus",
                    "excel_documents_milvus",
                    "json_documents_milvus",
                    "pdf_documents_milvus",
                    "txt_documents_milvus"
                ]
            )
            logger.info(f"服务初始化成功（可用集合数：{len(query_service.vector_stores)}）")

            # 构造过滤条件（从查询参数获取）
            custom_filter = {"file_name LIKE '%s'" % file_filter}

            # 调用查询
            result = await query_service.query_rag(
                user_question=question,
                user_filter=custom_filter
            )

            # 处理结果（保持不变）
            if result["status"] == "error":
                logger.info(f"查询失败：{result['error']['message']}")
                return ApiResponse(
                    status=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    message=result["error"]["message"],
                    data=result
                )

            logger.info(f"LLM响应：{result['response'][:50]}...")
            # 计算并打印执行时间
            elapsed_time = time.time() - start_time
            logger.info(f"函数执行耗时: {elapsed_time:.4f}秒")

            return ApiResponse(
                status=status.HTTP_200_OK,
                message="success",
                data=result
            )

        except RAGException as e:
            error_msg = f"服务初始化失败：[{e.code}] {e.message}"
            logger.error(error_msg, exc_info=True)
            return ApiResponse(
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
                message=error_msg
            )
        except Exception as e:
            error_msg = f"查询发生未知错误：{str(e)}"
            logger.error(error_msg, exc_info=True)
            return ApiResponse(
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
                message=error_msg
            )
        finally:
            if query_service:
                query_service.close()
    async def save_chunk(
            self,
            request: SystemQueryRequest
    ) -> ApiResponse:
        """
        核心逻辑：1.解析文档，切分数据块到milvus
        """
        try:
            # 1. 初始化服务
            logger.info("=" * 50)
            logger.info("开始初始化文档解析服务")
            parse_service = FilesParseService(default_collection=f"{request.question}_documents_milvus")
            logger.info("服务初始化成功")

            # 2. 处理docs目录下的所有文件
            logger.info("开始处理docs目录下的所有文件")
            dir_result = parse_service.process_docs_directory(f"./docs/contract/{request.question}")
            logger.info(f"\n【目录处理结果】\n{json.dumps(dir_result, ensure_ascii=False, indent=2)}")

            logger.info("\n" + "=" * 50)
            logger.info("所有文件处理完成")
            return ApiResponse(status=200,message="success",data=dir_result)

        except RAGException as e:
            logger.error(f"\n测试失败：[{e.code}] {e.message}", exc_info=True)
            exit(1)
        except Exception as e:
            logger.error(f"\n未知错误：{str(e)}", exc_info=True)
            exit(1)

